package leetcode.datastructure.heap;

import java.util.*;

/**
 * 此方法大数据量超时
 */
public class Program373KSmallestPairs {
    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        Comparator<List<Integer>> customComparator = (list1, list2) -> {
            int sum1 = list1.stream().mapToInt(Integer::intValue).sum();
            int sum2 = list2.stream().mapToInt(Integer::intValue).sum();
            return Integer.compare(sum1, sum2);
        };
        PriorityQueue<List<Integer>> heap = new PriorityQueue<>(k, customComparator);
        int maxIndex1 = Math.min(nums1.length , k);
        int maxIndex2 = Math.min(nums2.length , k);
        for (int i = 0; i < maxIndex1; i++) {
            for (int j = 0; j < maxIndex2; j++) {
                heap.offer(Arrays.asList(nums1[i], nums2[j]));
            }
        }
        List<List<Integer>> res = new ArrayList<>();
        int i = 0;
        while (!heap.isEmpty() && i < k) {
            res.add(heap.poll());
            i++;
        }
        return res;
    }

    public static void main(String[] args) {
        Program373KSmallestPairs myStack = new Program373KSmallestPairs();
        int[] nums1 = new int[]{1,7,11};
        int[] nums2 = new int[]{2,4,6};
        int k = 3;
        System.out.println(myStack.kSmallestPairs(nums1, nums2, k)); // 返回 2
    }
}
